{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'plt' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-1-f01034d01fe0>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     38\u001b[0m \u001b[1;31m# print(wf1d)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     39\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 40\u001b[1;33m \u001b[0mplt\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfigure\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     41\u001b[0m \u001b[0mplt\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mplot\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlinspace\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m,\u001b[0m\u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtime_series\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtime_series\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtime_series\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     42\u001b[0m \u001b[0mplt\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mplot\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlinspace\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m,\u001b[0m\u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mwf1d\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mwf1d\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mwf1d\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'plt' is not defined"
     ]
    }
   ],
   "source": [
    "def cross_correlation_wiener(ts, new_order, old_order):\n",
    "    n_taps = ts.size\n",
    "    n_size = n_taps if n_taps>new_order else new_order\n",
    "    out_data = []\n",
    "    for i in range(-n_taps,n_taps+1):\n",
    "        ts_sum = 0.0\n",
    "        for j in range(n_size):\n",
    "            if ((j+i)<n_taps and j<new_order and (j+i)>-1):\n",
    "                ts_sum += ts[j+i]\n",
    "        out_data.append(ts_sum)\n",
    "    out = out_data[-old_order-n_taps-1:-old_order-1]\n",
    "    return np.asarray(out)\n",
    "\n",
    "def wiener_filter_1d(t_series,order):\n",
    "    new_order = 2*order+1\n",
    "    t_2_series = pow(t_series,2)\n",
    "    t_lmean = cross_correlation_wiener(t_series, new_order, order)\n",
    "    lmean = t_lmean/new_order\n",
    "    t_lvar = cross_correlation_wiener(t_2_series, new_order, order)\n",
    "    lvar = t_lvar/new_order - pow(lmean,2)\n",
    "    est_n = sum(lvar)/lvar.size\n",
    "    res = []\n",
    "    for i in range(t_series.size):\n",
    "        if lvar[i] <est_n:\n",
    "            res.append(lmean[i])\n",
    "        else:\n",
    "            res.append((t_series[i]-lmean[i])*(1-(est_n/lvar[i]))+lmean[i])\n",
    "    return np.asarray(res)\n",
    "\n",
    "# 测试维纳滤波\n",
    "np.random.seed(1)\n",
    "t = np.linspace(0, 3*np.pi, 2048)\n",
    "time_series = np.sin(t) + 0.05*np.random.standard_normal(size=2048)\n",
    "order = 12\n",
    "wf1d = wiener_filter_1d(time_series, order)\n",
    "# print(wf1d)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(np.linspace(0,len(time_series), len(time_series)), time_series)\n",
    "plt.plot(np.linspace(0,len(wf1d), len(wf1d)), wf1d)\n",
    "plt.legend(['raw', 'filted'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
